import tensorflow as tf
import os

"""
这段代码的主要作用是：

加载原始的 SavedModel
创建一个带有标准输入签名的服务函数，假设模型输入是 RGB 图像
将这个带签名的函数与原始模型绑定
把带有签名的新模型保存到新目录
"""
# 原始模型路径
SAVED_MODEL_DIR = r"C:\Users\nanak\Desktop\weather-classification\saved_model"
# 新模型保存路径
NEW_SAVED_MODEL_DIR = r"C:\Users\nanak\Desktop\weather-classification\saved_model_with_signature"

# 确保新目录不存在，避免冲突
if os.path.exists(NEW_SAVED_MODEL_DIR):
    import shutil
    shutil.rmtree(NEW_SAVED_MODEL_DIR)

# 加载原始模型
try:
    loaded_model = tf.saved_model.load(SAVED_MODEL_DIR)
    print("✅ 成功加载原始模型")
except Exception as e:
    print(f"❌ 加载原始模型失败：{e}")
    exit(1)

# 创建一个带签名的函数包装器
@tf.function(input_signature=[
    tf.TensorSpec(shape=[None, None, None, 3],  # 假设输入是RGB图像，批次和尺寸可变
                 dtype=tf.float32, 
                 name='input_image')
])
def serving_function(input_image):
    # 尝试获取模型的默认调用方法
    if hasattr(loaded_model, '.__call__'):
        outputs = loaded_model(input_image)
    elif hasattr(loaded_model, 'call'):
        outputs = loaded_model.call(input_image)
    else:
        # 如果以上方法都不行，尝试直接调用模型
        outputs = loaded_model(input_image)
    
    # 如果输出是字典，直接返回；否则包装成字典
    if isinstance(outputs, dict):
        return outputs
    else:
        return {'outputs': outputs}

# 保存带有签名的新模型
tf.saved_model.save(
    loaded_model,
    NEW_SAVED_MODEL_DIR,
    signatures={'serving_default': serving_function}
)

print(f"✅ 已保存带签名的新模型到：{NEW_SAVED_MODEL_DIR}")

# 验证新模型
try:
    new_model = tf.saved_model.load(NEW_SAVED_MODEL_DIR)
    print("新模型包含的签名：", list(new_model.signatures.keys()))
    print("✅ 签名添加成功，可以用于TFLite转换了")
except Exception as e:
    print(f"❌ 验证新模型失败：{e}")
